Data Analysis Assistant

Introduction

One of the important aspects of data analysis is making outputs that are easy to understand and aesthetically pleasing. In my work, I refer to this as making graphs “pretty”. Today I want to test whether LLMs have a concept of “prettiness”. I’m interested to see how they behave when I ask them to modify graph code to make it more aesthetically pleasing.

For reference, this is what the original graph looks like:

import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, locations="Code",
                    color="Anxiety_percent",
                    hover_name="Entity", 
                    color_continuous_scale=px.colors.sequential.Plasma)
fig.show()

Test 1: GPT 4 Turbo

from openai import OpenAI
import os
import pandas as pd

from IPython.display import Markdown, display

from dotenv import load_dotenv
# Load API key
_ = load_dotenv()

client = OpenAI()

model = "gpt-4-turbo-2024-04-09"
system_role = '''You are a consultant skilled at making visually appealing and easy to understand
graphs to present to clients.'''
initial_prompt = '''Improve this code to make the graph better. Reply only with Python code.

Here is the code to improve: 
import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, locations="Code",
                    color="Anxiety_percent",
                    hover_name="Entity", 
                    color_continuous_scale=px.colors.sequential.Plasma)
fig.show()
'''

messages = [
    {"role": "system", "content": system_role} ,
    {"role": "user", "content": initial_prompt},
]

completion = client.chat.completions.create(
  model=model,        # the model to use 
  messages=messages,  # a list of dictionaries with the prompt
  n=3,                # how many responses to generate
  temperature=0.7     # temperature parameter
)

response = display(Markdown(f'-----\n#### Response {0} \n {resp.message.content} \n'))

---

Response 0

import pandas as pd
import plotly_express as px

# Loading data
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

# Creating a choropleth map with improved features
fig = px.choropleth(anx2010_df, locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity", 
                   title="Global Anxiety Rates in 2010",
                   labels={"Anxiety_percent": "Anxiety Rate (%)"},
                   color_continuous_scale=px.colors.sequential.Plasma,
                   projection="natural earth")

# Enhancing layout
fig.update_layout(
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   coloraxis_colorbar=dict(
       title="Anxiety Rate %",
       thicknessmode="pixels", thickness=15,
       lenmode="pixels", len=300
   )
)

# Showing the plot
fig.show()

Testing the code

import pandas as pd
import plotly_express as px

# Loading data
anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

# Creating a choropleth map with improved features
fig = px.choropleth(anx2010_df, locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity", 
                   title="Global Anxiety Rates in 2010",
                   labels={"Anxiety_percent": "Anxiety Rate (%)"},
                   color_continuous_scale=px.colors.sequential.Plasma,
                   projection="natural earth")

# Enhancing layout
fig.update_layout(
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   coloraxis_colorbar=dict(
       title="Anxiety Rate %",
       thicknessmode="pixels", thickness=15,
       lenmode="pixels", len=300
   )
)

# Showing the plot
fig.show()

Analysis

Here we can see that we have:

1. A title for the overall graph
2. A label for the anxiety rate legend
3. No more black box around the outside of the map

Test 2: Claude 3.5 Sonnet

from dotenv import load_dotenv
from anthropic import Anthropic
import base64
import httpx
from IPython.display import display, Image, Markdown

# load api keys
_ = load_dotenv()
client = Anthropic()

message = client.messages.create(
    model='claude-3-5-sonnet-20241022',
    max_tokens=1000,
    temperature=0.0,
    system='''You are a consultant skilled at making visually appealing and easy to understand
graphs to present to clients.''',
    messages=[
        {"role": "user", "content": '''Improve this code to make the graph better. Reply only with Python code.

Here is the code to improve: 
import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, locations="Code",
                    color="Anxiety_percent",
                    hover_name="Entity", 
                    color_continuous_scale=px.colors.sequential.Plasma)
fig.show()
'''}
    ]
)

response = display(Markdown(f'-----\n#### Response {0} \n {message.content[0].text} \n'))

---

Response 0

import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, 
                   locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity",
                   hover_data={"Anxiety_percent": ":.2f"},
                   color_continuous_scale="RdYlBu_r",
                   title="Global Anxiety Prevalence (2010)",
                   labels={"Anxiety_percent": "Anxiety Prevalence (%)",
                          "Code": "Country Code",
                          "Entity": "Country"})

fig.update_layout(
   title_x=0.5,
   title_font_size=24,
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   width=1200,
   height=700,
   margin=dict(t=50, l=0, r=0, b=0)
)

fig.update_coloraxes(
   colorbar_title_text="Prevalence (%)",
   colorbar_title_side="right",
   colorbar_thickness=15,
)

fig.show()

Testing the code

import pandas as pd
import plotly_express as px

anx2010_df = pd.read_csv('data/global_anxiety_2010.csv')

fig = px.choropleth(anx2010_df, 
                   locations="Code",
                   color="Anxiety_percent",
                   hover_name="Entity",
                   hover_data={"Anxiety_percent": ":.2f"},
                   color_continuous_scale="RdYlBu_r",
                   title="Global Anxiety Prevalence (2010)",
                   labels={"Anxiety_percent": "Anxiety Prevalence (%)",
                          "Code": "Country Code",
                          "Entity": "Country"})

fig.update_layout(
   title_x=0.5,
   title_font_size=24,
   geo=dict(
       showframe=False,
       showcoastlines=True,
       projection_type='equirectangular'
   ),
   width=1200,
   height=700,
   margin=dict(t=50, l=0, r=0, b=0)
)

fig.update_coloraxes(
   colorbar_title_text="Prevalence (%)",
   colorbar_title_side="right",
   colorbar_thickness=15,
)

fig.show() 

Analysis

Here we see:

1. A new color scheme
2. Updated titles

Conclusion

Both models did pretty well with this task. A lot of a graph’s “prettiness” comes down to personal preference, but asking LLMs for various options could definitely help in the brainstorming process.